package cas.ihep.spark.test;

import cas.ihep.fs.UnifiedBlock;
import cas.ihep.fs.UnifiedFile;
import cas.ihep.fs.UnifiedFileSystem;
import cas.ihep.util.MemoryBuffer;
import com.google.common.io.Closer;

import java.io.*;
import java.net.URI;

public class Benchmark {

    private static void test0(UnifiedFileSystem ufs,Closer closer, URI path) throws IOException {
        long bytes=0;
        int each;
        byte buf[]=new byte[8192];
        long startTime=System.nanoTime();
        BufferedInputStream bis=closer.register(new BufferedInputStream(ufs.openStream(path)));
        while((each=bis.read(buf))>0){
            bytes+=each;
        }
        startTime=System.nanoTime()-startTime;
        System.out.println("Throughput: "+bytes*1e9/startTime/1024/1024+"MB/s");
    }

    private static void test1(UnifiedFileSystem ufs,Closer closer,URI path)throws IOException{
        long bytes=0;
        int each;
        byte buf[]=new byte[8192];
        long startTime=System.nanoTime();
        UnifiedFile file=closer.register(ufs.open(path));
        UnifiedBlock block=file.block(0);
        BufferedInputStream bis=closer.register(new BufferedInputStream(new FileInputStream(block.absolutePath())));
        while((each=bis.read(buf))>0){
            bytes+=each;
        }
        startTime=System.nanoTime()-startTime;
        System.out.println("Throughput: "+bytes*1e9/startTime/1024/1024+"MB/s");
    }

    private static void test2(UnifiedFileSystem ufs,Closer closer,URI path)throws IOException{
        int each;
        byte buf[]=new byte[8192];
        int sz;
        long startTime=System.nanoTime();
        UnifiedFile file=closer.register(ufs.open(path));
        UnifiedBlock block=file.block(0);
        MemoryBuffer mbuf=closer.register(new MemoryBuffer(new File(block.absolutePath())));
        int capacity=mbuf.capacity();
        for(sz=0;sz<capacity;){
            each=Math.min(8192,capacity-sz);
            mbuf.get(buf,0,each);
            sz+=each;
        }
        startTime=System.nanoTime()-startTime;
        System.out.println("Throughput: "+sz*1e9/startTime/1024/1024+"MB/s");
    }

    private static final byte public_buf[]=new byte[]{(byte) 0xff, (byte) 0xfe, (byte) 0xfd, (byte) 0xfc, (byte) 0xfb, (byte) 0xfa, (byte) 0xf9, (byte) 0xf8};

    private static void test3(UnifiedFileSystem ufs,Closer closer,URI path,long bytes)throws IOException{
        long startTime=System.nanoTime();
        DataOutputStream dos=closer.register(new DataOutputStream(new BufferedOutputStream(ufs.createStream(path))));
        long written=0;
        for(;written<bytes;){
            int each=(int)(bytes-written);
            if(each>8){
                dos.writeLong(written);
                written+=8;
            }else{
                dos.write(public_buf,0,each);
                written+=each;
            }
        }
        startTime=System.nanoTime()-startTime;
        System.out.println("Throughput: "+bytes*1e9/startTime/1024/1024+"MB/s");
    }

    private static void test4(UnifiedFileSystem ufs,Closer closer,URI path,long bytes)throws IOException{
        long startTime=System.nanoTime();
        UnifiedFile file=closer.register(ufs.create(path));
        UnifiedBlock block=file.next();
        DataOutputStream dos=closer.register(new DataOutputStream(new BufferedOutputStream(new FileOutputStream(block.absolutePath()))));
        long written=0;
        for(;written<bytes;){
            int each=(int)(bytes-written);
            if(each>8){
                dos.writeLong(written);
                written+=8;
            }else{
                dos.write(public_buf,0,each);
                written+=each;
            }
        }
        startTime=System.nanoTime()-startTime;
        System.out.println("Throughput: "+bytes*1e9/startTime/1024/1024+"MB/s");
    }

    private static void test5(UnifiedFileSystem ufs,Closer closer,URI path,long bytes)throws IOException{
        long startTime=System.nanoTime();
        UnifiedFile file=closer.register(ufs.create(path));
        UnifiedBlock block=file.next();
        MemoryBuffer mbuf=closer.register(new MemoryBuffer(new File(block.absolutePath()),0,bytes));
        long written=0;
        for(;written<bytes;){
            int each=(int)(bytes-written);
            if(each>8){
                mbuf.putLong(written);
                written+=8;
            }else{
                mbuf.put(public_buf,0,each);
                written+=each;
            }
        }
        startTime=System.nanoTime()-startTime;
        System.out.println("Throughput: "+bytes*1e9/startTime/1024/1024+"MB/s");
    }

    public static void main(String args[])throws Exception{
        try(Closer closer=Closer.create()){
            UnifiedFileSystem ufs=closer.register(UnifiedFileSystem.get(UnifiedFileSystem.newConfiguration(System.getenv("HADOOP_HOME"))));
            int type=Integer.parseInt(args[0]);
            URI uri=new URI("alluxio://hadoop06.ihep.ac.cn:19998"+args[1]);
            switch(type){
                case 0:
                    test0(ufs,closer,uri);
                    break;
                case 1:
                    test1(ufs,closer,uri);
                    break;
                case 2:
                    test2(ufs,closer,uri);
                    break;
                case 3:
                    test3(ufs,closer,uri,Long.parseLong(args[2])*1024*1024);
                    break;
                case 4:
                    test4(ufs,closer,uri,Long.parseLong(args[2])*1024*1024);
                    break;
                default:
                    test5(ufs,closer,uri,Long.parseLong(args[2])*1024*1024);
            }
        }
    }
}
